﻿#light

module DwarfFSharpMpi

open System
open System.Collections.Generic
open System.Runtime.InteropServices
open MPI
open System.Threading
open System.Numerics
open DwarfConfigurator

[<DllImport(@"CliTools.dll")>]
extern void start()

[<DllImport(@"CliTools.dll")>]
extern double finish()

/// <summary>
/// 1-dimension Fast Fourier Algorithm
/// </summary>
let fft length =
    // Create the array of phase weights
    let W = Array.zeroCreate length
    let factor = -2.0 * Math.PI / (float length);
    for i = 0 to length - 1 do
        let real = cos((float i) * factor)
        let imag = sin((float i) * factor)
        W.[i] <- new Complex(real,imag)
    // Now define a function that evaluates the Fourier transform
    let FFT (line : Complex array) = 
        //variables for the FFT
        let length = line.Length
        let bitwidth =  int (Math.Log((float)length, (float)2))

        //reorder for butterflies
        let firstBit = length >>> 1
        let mutable index = 0
        let mutable reverseIndex = 0
        
        for index = 0 to length - 1 do
            if (reverseIndex < index) then
                let w = line.[index]
                line.[index] <- line.[reverseIndex]
                line.[reverseIndex] <- w
                
            // Add 1 to reverse index, going from left to right:
            let mutable carry = firstBit
            // If adding 1 changed the bit to 0, we have a carry:
            
            reverseIndex <- (reverseIndex ^^^ carry)
            if ((length - 1) <> index) then
                while ((reverseIndex &&& carry) = 0 ) do
                    carry <- (carry >>> 1)
                    reverseIndex <- (reverseIndex ^^^ carry)

        //combine and untangle applying phase weights
        //using successive butterflies of 2,4,...,bitwidth
        for i = 1 to bitwidth do
            let halfStep = 1 <<< (i - 1)
            let step = halfStep + halfStep
            for j = 0 to halfStep - 1 do
                let tempShift = j <<< (bitwidth - i)
                let w = W.[tempShift]
                // We could use a for loop here:
                //     for k in j .. halfStep * 2 .. length - 1 do
                // but that is slow in the current F# release.
                let mutable k = j
                while (k < length) do
                    let u = line.[k + halfStep]
                    let v = line.[k]
                    let t = w * u
                    line.[k + halfStep] <- v - t
                    line.[k] <- v + t
                    k <- k + step
    // Return the function we just defined
    FFT

/// <summary>
/// Transposition of matrix
/// </summary>
let transpose (x : Complex array array ) =
    let size = x.GetLength(0)
    for i = 0 to size - 1 do
        for j = i+1 to size - 1 do
            let t = x.[i].[j]
            x.[i].[j] <- x.[j].[i]
            x.[j].[i] <- t


/// <summary>
/// Send row and index of row to dest
/// </summary>
/// <param name="tRow">array of type Complex[] to send</param>
/// <param name="dest">MPI rank of process to receive row</param>
/// <param name="index">Index of row in original matrix</param>         
let SendRow(row : Complex[], dest : int, index : int) =
    Intracommunicator.world.Send<Complex>(row,dest,index)

/// <summary>
/// Receive a row from requested source
/// </summary>
/// <param name="source">MPI rank to receive row from.  To receive from any, Use Communicator.anySource</param>
/// <param name="row">Array to receive the row</param>
/// <param name="index">Set on return to the row index of the received row in the orginalmatrix</param>
/// <returns>MPI.CompletedStatus of received row</returns>
let ReceiveRow(source: int, row : Complex[], index : int) : CompletedStatus =
    let complStatus = Intracommunicator.world.Receive<Complex>(source,index,ref row) : CompletedStatus
    complStatus


/// <summary>
/// Evaluate an action on every row in matrix
/// </summary>
/// <remarks>
/// The action will be evaluated once for each row 
/// with the row passed to the action as the sole argument, 
/// The action is allowed to change the contents of the row
/// </remarks>
/// <typeparam name="T">Element type of matrix</typeparam>
/// <param name="matrix">matrix to operate on</param>
/// <param name="Row_Action">Single argument action to perform</param>
let EvaluateRowsUsingMPI(model : DataModel) = 
    //set bounds of matrix for worker processes
    let size = (int) model.length;
    
    let myFft = fft size
    let MpiSize = Communicator.world.Size
    let colLength = (int) model.length
    let rowLength = (int) model.length

    //Set All Done tag for worker processes
    DataModel.AllDoneTag <- colLength

    if (model.isRootThread) then 
        //Root Process
        let mutable rowsSent = 0
        let mutable rowsProcessed = 0

        //Give each worker process some rows to start with
        for i = 1 to MpiSize - 1 do
            SendRow(model.complexArray.[rowsSent % colLength], i, rowsSent)
            rowsSent <- rowsSent + 1

        //Loop until all rows sent and all transformed rows received
        let mutable goToNextItaration = false
        while (rowsProcessed < colLength) do
            goToNextItaration <- false
            //Check if any results sent back from any worker
            let probeStatus = Communicator.world.ImmediateProbe(Communicator.anySource, Communicator.anyTag)
            if (probeStatus = null) then
                //nothing to receive, do we have any rows left to do
                if (rowsSent < colLength) then
                    //yes, transform next row in root process
                    myFft model.complexArray.[rowsSent]
                    rowsSent <- rowsSent + 1
                    rowsProcessed <- rowsProcessed + 1
                else 
                    Thread.Sleep(0)  //we have nothing to do
                goToNextItaration <- true
            
            if (goToNextItaration = false) then 
                //receive row, waiting if necessary                
                let recvStatus = ReceiveRow(probeStatus.Source,model.complexArray.[probeStatus.Tag], probeStatus.Tag)
                let source = recvStatus.Source
                let recvIndex = recvStatus.Tag
                rowsProcessed <- rowsProcessed + 1            

                //Send next row (if any) to process  that completed.
                if (rowsSent < colLength) then
                    SendRow(model.complexArray.[rowsSent], source, rowsSent)
                    rowsSent <- rowsSent + 1

        //Tell workers all done
        let complex = new Complex()
        let row = Array.create rowLength complex

        for i = 1 to MpiSize - 1 do
            SendRow(row, i, DataModel.AllDoneTag)
            
    else
        //Worker Processes

        //Process rows until told to stop.
        let complex = new Complex()
        let row = Array.create rowLength complex
        let mutable notBreak = true
        while (notBreak) do
            //Receive row, waiting if necessary.
            let mutable rowIndex = 0
             
            let recvStatus = ReceiveRow(DataModel.MpiRoot, row, Communicator.anyTag)
            rowIndex <- recvStatus.Tag
            //are we done?
            if (rowIndex = DataModel.AllDoneTag) then
                //end of while cycle
                notBreak <- false
            else 
                //Transform the row 
                myFft row

                // send the transformed row back 
                SendRow(row, DataModel.MpiRoot, rowIndex)

let processSpectralMethod(model : DataModel) : DataModel =  
    // Transform the rows 
    EvaluateRowsUsingMPI(model)

    // Transform the columns
    if (model.isRootThread) then transpose model.complexArray 

    // Transform the rows 
    EvaluateRowsUsingMPI(model)
        
    // Transform the columns
    if (model.isRootThread) then transpose model.complexArray 
    
    model

let main() =
    let args = System.Environment.GetCommandLineArgs()
    using(new Environment(ref args))(fun _->
        let conf = new Configurator()
        let mutable model = conf.configure(args,true)
        if(model.isRootThread) then
            model.DWARF_NAME <- "Spectral Method, fsharp Mpi kernel."
            model.DWARF_NAME_OUTPUT <- "Dwarf.FSharp.Mpi"
            start()
            
        model <- processSpectralMethod(model)
        if(model.isRootThread) then
            let time = finish()
            conf.Finish(model,time)
    )
  
main()   

    